import pickle as pkl
import pandas as pd
import random
import numpy as np
import time

from bert_serving.client import BertClient
bc = BertClient(ip='localhost',check_version=False, check_length=False)

with open('./Data/Processed/word_raw.pkl', 'rb') as f:
    word_raw = pkl.load(f)
print(len(word_raw))

word_emb0 = []
timer = 0
for row in word_raw:
    if (timer % 1000 == 0):
        print(timer, "iterations")
    timer += 1
    word_emb0.append(bc.encode(row[:-1]))

word_emb1 = []
timer = 0
for row in word_raw:
    if (timer % 1000 == 0):
        print(timer, "iterations")
    timer += 1
    if row[-1] == "":
        word_emb1.append(None)
    else:
        word_emb1.append(bc.encode(row[-1:]))

cnt = 0
mean = np.zeros_like(word_emb1[0])
for x in word_emb1:
    if x is not None:
        cnt += 1
        mean += x
mean /= cnt
print((mean*mean).sum(axis=1))

for i in range(len(word_emb1)):
    if word_emb1[i] is None:
        word_emb1[i] = mean

word_emb = []
for embed0, embed1 in zip(word_emb0, word_emb1):
    word_emb.append(np.concatenate([embed0, embed1], axis=0).reshape((-1, )))
word_emb_tensor = np.stack(word_emb, axis=0)
print(word_emb_tensor.shape)

with open('./Data/Processed/word_emb_tensor.pkl', 'wb') as f:
    pkl.dump(word_emb_tensor, f, pkl.HIGHEST_PROTOCOL)

print("Done")
